#!/usr/bin/env python

# https://zhajiman.github.io/post/chromaticity_diagram/
# http://www.cvrl.org/

import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon
import numpy as np
from colour import XYZ_to_xyY, xyY_to_XYZ, RGB_to_XYZ, XYZ_to_sRGB
import cv2
# pip install colour-science --break-system=packages
def xyY2RGB(x,y):
    x, y = np.meshgrid(x, y)

    Y = np.ones_like(y)

    Y_y = Y / y
    X = x * Y_y
    Z = (1 - x - y) * Y_y
    XYZ = np.dstack((X, Y, Z))

    M = np.array([
        [+3.2406, -1.5372, -0.4986],
        [-0.9689, +1.8758, +0.0415],
        [+0.0557, -0.2040, +1.0570]
    ])
    RGB = np.tensordot(XYZ, M, (-1, 1))
    RGB = RGB.clip(0, None)

    RGB -= RGB.min(axis=2, keepdims=True).clip(None, 0)
    RGB /= RGB.max(axis=2, keepdims=True)

    mask = RGB > 0.0031308
    RGB[~mask] *= 12.92
    RGB[mask] = 1.055 * RGB[mask]**(1 / 2.4) - 0.055
    return RGB



fig, ax = plt.subplots()
ax.set_aspect('equal')
ax.set_xlim(-0.1, 0.8)
ax.set_ylim(-0.1, 0.9)

cc = pd.read_csv('./cc2012xyz2_fine_5dp.csv', index_col=0)
patch = Polygon(
    xy=cc[['x', 'y']],
    transform=ax.transData,
    ec='k', fc='none', lw=1
)
ax.add_patch(patch)

N = 256
x = np.linspace(-0.1, 1, N)
y = np.linspace(-0.1, 1, N).clip(1e-3, 1)

RGB = xyY2RGB(x,y)
# x, y = np.meshgrid(x, y)

# Y = np.ones_like(y)
# Y_y = Y / y
# X = x * Y_y
# Z = (1 - x - y) * Y_y
# XYZ = np.dstack((X, Y, Z))

# M = np.array([
#     [+3.2406, -1.5372, -0.4986],
#     [-0.9689, +1.8758, +0.0415],
#     [+0.0557, -0.2040, +1.0570]
# ])

# RGB = np.tensordot(XYZ, M, (-1, 1))
# RGB = RGB.clip(0, None)
# RGB -= RGB.min(axis=2, keepdims=True).clip(None, 0)
# RGB /= RGB.max(axis=2, keepdims=True)
# mask = RGB > 0.0031308
# RGB[~mask] *= 12.92
# RGB[mask] = 1.055 * RGB[mask]**(1 / 2.4) - 0.055

ax.imshow(
    RGB,
    origin='lower',
    extent=[0, 1, 0, 1],
    interpolation='bilinear',
    clip_path=patch
)


x, y = cc['x'], cc['y']
xy = np.column_stack((x, y))
dxdy = np.zeros_like(xy)
dxdy[0] = xy[1] - xy[0]
dxdy[-1] = xy[-1] - xy[-2]
dxdy[1:-1] = xy[2:] - xy[:-2]
dx, dy = dxdy[:, 0], dxdy[:, 1]
dl = np.hypot(dx, dy)
dl[dl <= 0] = np.nan
cos = -dy / dl
sin = dx / dl

cs = pd.DataFrame(
    data=np.column_stack((cos, sin)),
    index=cc.index,
    columns=['cos', 'sin']
)
cs.loc[(cs.index < 430) | (cs.index > 660)] = np.nan
cs = cs.ffill().bfill()
cos, sin = cs['cos'], cs['sin']

tick_len = 0.03
tick_df = pd.DataFrame({
    'x0': x,
    'y0': y,
    'x1': x + tick_len * cos,
    'y1': y + tick_len * sin,
})

ticks = [390, *range(460, 601, 10), 620, 700]
tick_df = tick_df.loc[ticks]

for row in tick_df.itertuples():
    ax.plot(
        [row.x0, row.x1],
        [row.y0, row.y1],
        c='k', lw=0.6
    )
    if(row.Index < 510) :
        ax.text(row.x1,row.y1,'%3d'%(row.Index),ha='right',va='center')
    elif (row.Index > 540) :
        ax.text(row.x1,row.y1,'%3d'%(row.Index),ha='left',va='center')
    else :
        ax.text(row.x1,row.y1,'%3d'%(row.Index),ha='center',va='bottom')



sRGB_xy = [(0.64,0.33),(0.3,0.6),(0.15,0.06)]
ax.add_patch(Polygon(sRGB_xy,closed=True,fill=None,edgecolor='black'))

#led_xy = [(0.6801,0.3197),(0.2658,0.7243),(0.1510,0.0227)]

led_xy = [(0.7006,0.2993),(0.2658,0.7243),(0.1440,0.0297)]
ax.add_patch(Polygon(led_xy,closed=True,fill=None,edgecolor='black'))



def sRGB_to_linear_rgb(sRGB):
    linear_rgb = [pow((c + 0.055) / 1.055, 2.4) if c > 0.04045 else c / 12.92 for c in sRGB]
    return linear_rgb

def generate_rgb_grid(start, stop, steps):
    r = np.linspace(start[0], stop[0], steps)
    g = np.linspace(start[1], stop[1], steps)
    b = np.linspace(start[2], stop[2], steps)
    rgb = np.dstack((r, g, b))
    return [tuple(i) for i in rgb[0]]


def RGB2xyY(rgb):
#    lrgb = sRGB_to_linear_rgb(rgb)
    XYZ = RGB_to_XYZ(rgb, 'sRGB')
    Yxy = XYZ_to_xyY(XYZ)
    return (Yxy[0],Yxy[1])


RGB= generate_rgb_grid([0.001,0.001,0.001],[1.0,1.0,1.0],255)

xy = [RGB2xyY(c) for c in RGB]

ax.add_patch(Polygon(xy,closed=False,fill=None,edgecolor='black'))


def xy_to_RGB(xy, Y=1):
    print(xy[0],xy[1])
    XYZ = xyY_to_XYZ(np.array([xy[0], xy[1], Y]))
    linear_rgb = XYZ_to_sRGB(XYZ)
    # linear_rgb = [pow((c + 0.055) / 1.055, 2.4) if c > 0.04045 else c / 12.92 for c in linear_rgb]
    # linear_rgb = [(1.055 * pow(c, -2.4) - 0.055) if c > 0.0031308 else c * 12.92 for c in linear_rgb]
    # rgb = np.clip(linear_rgb, 0, 1)  # 确保数值在0到1之间
    return tuple(linear_rgb)
YY=[18.0, 74.5, 12.5]
for t in led_xy:
    rgb_color = xy_to_RGB(t,YY[0]/100.0)
    print(f"xy: {t} -> RGB: {rgb_color}")

print("---------------------")

for t in sRGB_xy:
    rgb_color = xy_to_RGB(t)
    print(f"xy: {t} -> RGB: {rgb_color}")

src_pts=np.float32(sRGB_xy)
dst_pts=np.float32(led_xy)

M = cv2.getAffineTransform(src_pts,dst_pts)
print(M)

x = [0.64,0.33]

src_point = np.array(x,dtype='float32')
src_point_reshaped = src_point.reshape(-1, 1, 2)

dst_point_reshaped = cv2.transform(src_point_reshaped, M)

print(dst_point_reshaped)

plt.show()
